from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
import pymysql
from tqdm import tqdm

device = 'cuda'
batch_size = 1000  # 每批处理的记录数
# 数据库连接配置

semantic_cls = pipeline(Tasks.text_classification,
                        r'F:\model\emotion\nlp_structbert_sentiment-classification_chinese-base', device=device)

def emotion_get(comment):
    emotion_pro = semantic_cls(input=comment)
    if emotion_pro['scores'][0] > 0.88:
        emotion = emotion_pro['labels'][0]
    else:
        emotion = '中性'
    return emotion

if __name__ == '__main__':
    DB_CONFIG = {
        'host': 'rm-2zea30h4sh8g15zd1ho.mysql.rds.aliyuncs.com',
        'port': 3306,
        'user': "root",
        'password': 'Ds2024@()833429',
        'database': "douyinpinglun"
    }
    mydb = pymysql.connect(**DB_CONFIG)
    cursor = mydb.cursor()
    try:
        offset = 0
        while True:
            # 从表中选择 id 和 comment_text，带有 LIMIT 和 OFFSET
            query = """
                select id,评论内容 from 视频评论表 where 情感分类 is null
            """
            cursor.execute(query)
            rows = cursor.fetchall()
            if not rows:
                break  # 如果没有更多记录，退出循环
            # 更新每条记录的情感值
            for row in tqdm(rows, desc=f"Processing batch starting at offset {offset}"):
                record_id = row[0]
                comment_text = row[1]

                # 获取情感值
                emotion = emotion_get(comment_text)

                # 更新数据库中的 emotion 字段
                cursor.execute("""
                    update 视频评论表 set 情感分类=%s where id=%s
                """, (emotion, record_id))
            # 提交更改
            mydb.commit()
            # 更新偏移量
            offset += batch_size
            print(f'{offset}条数据处理完毕')

    finally:
        # 关闭游标和连接
        cursor.close()
        mydb.close()

    print("情感分析完成，并且数据库更新成功。")
